import os

from tqdm import tqdm

from ShieldNet import ModelTrainer, BENIGN_LABEL, ClassifyName, ModelPlotSaver, ModelCreatorConfig, \
    DataLoader, TransformerDDoSDetector


def train_model(data_path: str, train_name: str):
    # 实例化模型训练器
    trainer = ModelTrainer(
        device='cuda',
        seed=-1,
        epochs=250,
        batch_size=512,
        lr=1e-4,
        window_size=100,
        step=5,
        is_debug=True,
        model_plot_saver=ModelPlotSaver(
            plot_train_save_path=os.path.join(PLOT_SAVE_PATH, train_name, 'train.png'),
            plot_cm_save_path=os.path.join(PLOT_SAVE_PATH, train_name, 'cm.png'),
            plot_pr_save_path=os.path.join(PLOT_SAVE_PATH, train_name, 'pr.png'),
        ),
        model_creator_config=ModelCreatorConfig(
            model_path=MODEL_SAVE_PATH_,
            model_name=f"Model_{train_name}",
            suffix="model",
            pth_path=os.path.join(MODEL_SAVE_PATH, train_name, "model.pth"),
            static_pth_path=os.path.join(MODEL_SAVE_PATH, train_name, "static", "model_static.pth"),
            pkl_path=os.path.join(SCALER_SAVE_PATH, train_name, "scaler.pkl"),
        ),
        classify_name=ClassifyName.LABEL.value,
        classify_loader=[BENIGN_LABEL, train_name],
        data_loader=DataLoader(chunk_size=1000000),
        detector=TransformerDDoSDetector
    )
    # 开始训练
    (trainer
     .load_data(data_path)
     .deal_with()
     .divide_dataset()
     .init_model()
     .train()
     .show_train_plot()
     .model_evaluate()
     .plot_pr_curve()
     .clear_model())

if __name__ == '__main__':
    SCALER_SAVE_PATH = r'E:\C4\Ai\Train\Skl'
    MODEL_SAVE_PATH = r'E:\C4\Ai\Train\Pth'
    PLOT_SAVE_PATH = r'E:\C4\Ai\Train\Plot'
    JSON_SAVE_PATH = r'E:\C4\Ai\Train\Json'
    MODEL_SAVE_PATH_ = r'E:\C4\Ai\Train\Model'

    # 训练内容：
    train_list = [
        (r"E:\C4\Ai\data\Syn.csv", "Syn"),
        (r"E:\C4\Ai\data\TFTP.csv", "TFTP"),
        (r"E:\C4\Ai\data\UDPLag.csv", "UDPLag"),
        (r"E:\C4\Ai\data\DrDoS_DNS.csv", "DrDoS_DNS"),
        (r"E:\C4\Ai\data\DrDoS_NTP.csv", "DrDoS_NTP"),
        (r"E:\C4\Ai\data\DrDoS_LDAP.csv", "DrDoS_LDAP"),
        (r"E:\C4\Ai\data\DrDoS_SSDP.csv", "DrDoS_SSDP"),
        (r"E:\C4\Ai\data\DrDoS_SNMP.csv", "DrDoS_SNMP"),
        (r"E:\C4\Ai\data\DrDoS_MSSQL.csv", "DrDoS_MSSQL"),
        (r"E:\C4\Ai\data\DrDoS_NetBIOS.csv", "DrDoS_NetBIOS"),
        (r"E:\C4\Ai\data\DrDoS_UDP.csv", "DrDoS_UDP"),
    ]

    for content in train_list:
        train_model(*content)


